from .capture_base import CaptureBase
import numpy as np


class MSSCapture(CaptureBase):
    def __init__(self, monitor=1, width=320, height=320):
        """
        :param monitor: 显示器编号（1 表示第一个显示器）
        :param width: 截图宽度
        :param height: 截图高度
        """
        self.monitor_index = monitor
        self.width = width
        self.height = height
        self.sct = None
        self.region = None

    def init_in_thread(self):
        """在目标线程里初始化 mss"""
        import mss
        self.sct = mss.mss()
        self._update_region()

    def _update_region(self):
        """根据当前宽高重新计算截图区域"""
        if self.sct is None:
            raise RuntimeError("MSS 尚未初始化，请先调用 init_in_thread()")

        monitors = self.sct.monitors
        if self.monitor_index < 1 or self.monitor_index >= len(monitors):
            self.monitor_index = 1
        m = monitors[self.monitor_index]

        left = m["left"] + (m["width"] - self.width) // 2
        top = m["top"] + (m["height"] - self.height) // 2
        self.region = {
            "left": left,
            "top": top,
            "width": self.width,
            "height": self.height,
        }

    def update_size(self, width, height):
        """热更新截图区域大小"""
        self.width = width
        self.height = height
        self._update_region()
        print(f"🔄 截图大小已更新: {self.width}x{self.height}")

    def grab_frame(self):
        """抓取一帧 (H, W, 3)，BGR"""
        if self.sct is None or self.region is None:
            raise RuntimeError("请先调用 init_in_thread()")
        frame = self.sct.grab(self.region)
        return np.array(frame)[:, :, :3]
